#pragma once

#include "viml/ast.h"
#include <sstream>

namespace viml {

// AST Pretty Printer - 提供格式化的AST输出
class ASTPrettyPrinter : public ASTVisitor {
private:
    std::ostringstream output;
    int indent_level;
    
    void print_indent() {
        for (int i = 0; i < indent_level; ++i) {
            output << "  ";
        }
    }
    
    void increase_indent() { indent_level++; }
    void decrease_indent() { indent_level--; }
    
public:
    ASTPrettyPrinter() : indent_level(0) {}
    
    std::string get_output() const { return output.str(); }
    
    void visit(NumberLiteral& node) override {
        output << "NumberLiteral(" << node.value << ")";
    }
    
    void visit(StringLiteral& node) override {
        output << "StringLiteral(\"" << node.value << "\")";
    }
    
    void visit(Identifier& node) override {
        output << "Identifier(" << node.name << ")";
    }
    
    void visit(SpecialConstantNode& node) override {
        output << "SpecialConstantNode(" << node.constant_name << ")";
    }
    
    void visit(BinaryExpression& node) override {
        output << "BinaryExpression(\n";
        increase_indent();
        
        print_indent();
        output << "operator: " << BinaryExpression::operator_to_string(node.op) << "\n";
        
        print_indent();
        output << "left: ";
        node.left->accept(*this);
        output << "\n";
        
        print_indent();
        output << "right: ";
        node.right->accept(*this);
        output << "\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(UnaryExpression& node) override {
        output << "UnaryExpression(\n";
        increase_indent();
        
        print_indent();
        output << "operator: " << UnaryExpression::operator_to_string(node.op) << "\n";
        
        print_indent();
        output << "operand: ";
        node.operand->accept(*this);
        output << "\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(CallExpression& node) override {
        output << "CallExpression(\n";
        increase_indent();
        
        print_indent();
        output << "callee: ";
        node.callee->accept(*this);
        output << "\n";
        
        print_indent();
        output << "arguments: [\n";
        increase_indent();
        for (size_t i = 0; i < node.arguments.size(); ++i) {
            print_indent();
            node.arguments[i]->accept(*this);
            if (i < node.arguments.size() - 1) output << ",";
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "]\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(ListLiteral& node) override {
        output << "ListLiteral([\n";
        increase_indent();
        for (size_t i = 0; i < node.elements.size(); ++i) {
            print_indent();
            node.elements[i]->accept(*this);
            if (i < node.elements.size() - 1) output << ",";
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "])";
    }
    
    void visit(DictLiteral& node) override {
        output << "DictLiteral({\n";
        increase_indent();
        for (size_t i = 0; i < node.pairs.size(); ++i) {
            print_indent();
            node.pairs[i].first->accept(*this);
            output << ": ";
            node.pairs[i].second->accept(*this);
            if (i < node.pairs.size() - 1) output << ",";
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "})";
    }
    
    void visit(IndexExpression& node) override {
        output << "IndexExpression(\n";
        increase_indent();
        
        print_indent();
        output << "object: ";
        node.object->accept(*this);
        output << "\n";
        
        print_indent();
        output << "index: ";
        node.index->accept(*this);
        output << "\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(LetStatement& node) override {
        output << "LetStatement(\n";
        increase_indent();
        
        print_indent();
        output << "name: " << node.name << "\n";
        
        print_indent();
        output << "value: ";
        node.value->accept(*this);
        output << "\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(EchoStatement& node) override {
        output << "EchoStatement(\n";
        increase_indent();
        
        print_indent();
        output << "expression: ";
        node.expression->accept(*this);
        output << "\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(IfStatement& node) override {
        output << "IfStatement(\n";
        increase_indent();
        
        print_indent();
        output << "condition: ";
        node.condition->accept(*this);
        output << "\n";
        
        print_indent();
        output << "then_body: [\n";
        increase_indent();
        for (auto& stmt : node.then_body) {
            print_indent();
            stmt->accept(*this);
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "]\n";
        
        if (!node.else_body.empty()) {
            print_indent();
            output << "else_body: [\n";
            increase_indent();
            for (auto& stmt : node.else_body) {
                print_indent();
                stmt->accept(*this);
                output << "\n";
            }
            decrease_indent();
            print_indent();
            output << "]\n";
        }
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(WhileStatement& node) override {
        output << "WhileStatement(\n";
        increase_indent();
        
        print_indent();
        output << "condition: ";
        node.condition->accept(*this);
        output << "\n";
        
        print_indent();
        output << "body: [\n";
        increase_indent();
        for (auto& stmt : node.body) {
            print_indent();
            stmt->accept(*this);
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "]\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(ForStatement& node) override {
        output << "ForStatement(\n";
        increase_indent();
        
        print_indent();
        output << "variable: " << node.variable << "\n";
        
        print_indent();
        output << "iterable: ";
        node.iterable->accept(*this);
        output << "\n";
        
        print_indent();
        output << "body: [\n";
        increase_indent();
        for (auto& stmt : node.body) {
            print_indent();
            stmt->accept(*this);
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "]\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(FunctionStatement& node) override {
        output << "FunctionStatement(\n";
        increase_indent();
        
        print_indent();
        output << "name: " << node.name << "\n";
        
        print_indent();
        output << "parameters: [";
        for (size_t i = 0; i < node.parameters.size(); ++i) {
            output << node.parameters[i];
            if (i < node.parameters.size() - 1) output << ", ";
        }
        output << "]\n";
        
        print_indent();
        output << "body: [\n";
        increase_indent();
        for (auto& stmt : node.body) {
            print_indent();
            stmt->accept(*this);
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "]\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(ReturnStatement& node) override {
        output << "ReturnStatement(\n";
        increase_indent();
        
        if (node.value) {
            print_indent();
            output << "value: ";
            node.value->accept(*this);
            output << "\n";
        } else {
            print_indent();
            output << "value: null\n";
        }
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(CallStatement& node) override {
        output << "CallStatement(\n";
        increase_indent();
        
        print_indent();
        output << "call_expr: ";
        node.call_expr->accept(*this);
        output << "\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
    
    void visit(Program& node) override {
        output << "Program(\n";
        increase_indent();
        
        print_indent();
        output << "statements: [\n";
        increase_indent();
        for (size_t i = 0; i < node.statements.size(); ++i) {
            print_indent();
            node.statements[i]->accept(*this);
            if (i < node.statements.size() - 1) output << ",";
            output << "\n";
        }
        decrease_indent();
        print_indent();
        output << "]\n";
        
        decrease_indent();
        print_indent();
        output << ")";
    }
};

} // namespace viml